#include "X264VideoEncoder.h"
#include "common.h"
#include <stdio.h>
#include <string.h>
#include <assert.h>
#include <Windows.h>

static uint8_t startCode[4] = {0, 0, 0, 1};

CX264VideoEncoder::CX264VideoEncoder()
: m_x264(NULL)
, m_param(NULL)
, m_pic(NULL)
, m_spsLen(0)
, m_ppsLen(0)
, m_frames(0)
#if LOG_INFO
, m_logFile(NULL)
, m_encodedTime(0.0)
#endif
{
	m_param = new x264_param_t;
	m_pic = new x264_picture_t;
}

CX264VideoEncoder::~CX264VideoEncoder()
{
	if (m_pic)
	{
		delete m_pic;
		m_pic = NULL;
	}
	if (m_param)
	{
		delete m_param;
		m_param = NULL;
	}
#if LOG_INFO
	if (m_logFile)
	{
		fclose(m_logFile);
		m_logFile = NULL;
	}
#endif
}

bool CX264VideoEncoder::Open( int aWidth, int aHeight, int aAnnexB, 
								uint8_t * aSPSBuffer, int * aSPSLen, 
								uint8_t * aPPSBuffer, int * aPPSLen )
{
	x264_param_default(m_param);
	LoadConfig("./encoder.conf");

	m_param->i_width = aWidth;
	m_param->i_height = aHeight;

#if LOG_INFO
	m_encodedTime = 0.0;
#endif
	m_frames = 0;
	m_x264 = x264_encoder_open(m_param);
	if (!m_x264)
		return false;

	x264_picture_alloc(m_pic, X264_CSP_I420, m_param->i_width, m_param->i_height);

	UpdatePS();

	if (NULL == aSPSBuffer && NULL == aSPSLen &&
		NULL == aPPSBuffer && NULL == aPPSLen)
		return true;

	GetSPS(aSPSBuffer, aSPSLen, aAnnexB);
	GetPPS(aPPSBuffer, aPPSLen, aAnnexB);

	return true;
}

void CX264VideoEncoder::Close()
{
#if LOG_INFO
	if (m_logFile)
	{
		#define SUM3(p) (p[SLICE_TYPE_I] + p[SLICE_TYPE_P] + p[SLICE_TYPE_B])

		x264_t * h = m_x264;
		const int i_count = h->stat.i_slice_count[SLICE_TYPE_I] +
			h->stat.i_slice_count[SLICE_TYPE_P] +
			h->stat.i_slice_count[SLICE_TYPE_B];
		float fps = (float) h->param.i_fps_num / h->param.i_fps_den;
		float f_bitrate = fps * SUM3(h->stat.i_slice_size) / i_count / 125;

		if(h->param.analyse.b_ssim)
		{
			fprintf(m_logFile, "SSIM Mean Y:%.7f\n", SUM3( h->stat.f_ssim_mean_y ) / i_count);
		}
		if(h->param.analyse.b_psnr)
		{
			fprintf(m_logFile, "PSNR Mean Y:%6.3f\n", SUM3( h->stat.f_psnr_mean_y ) / i_count);
		}
		fprintf(m_logFile, "encoded %d frames, %.2f fps, %.2f kb/s\n", 
			m_frames,  m_frames / m_encodedTime, f_bitrate);

		fclose(m_logFile);
		m_logFile = NULL;

		#undef SUM3
	}
#endif

	if (m_pic)
		x264_picture_clean(m_pic);
	if (m_x264)
		x264_encoder_close(m_x264);
}

bool CX264VideoEncoder::Encode( uint8_t * inBuffer, uint8_t * outBuffer, int * outLen, int * frameType, int bAnnexB )
{
#if LOG_INFO
	LARGE_INTEGER litmp; 
	LONGLONG start, end;
	double freq, time; 
	QueryPerformanceFrequency(&litmp);
	freq = (double)litmp.QuadPart;
	QueryPerformanceCounter(&litmp);
	start = litmp.QuadPart;
#endif

	assert(inBuffer);
	assert(outBuffer);
	assert(outLen && *outLen > 0);
	assert(m_pic);
	assert(m_x264);

	int imgPixels = m_param->i_width * m_param->i_height;
	memcpy(m_pic->img.plane[0], inBuffer, imgPixels);
	memcpy(m_pic->img.plane[1], inBuffer + imgPixels, (imgPixels >> 2));
	memcpy(m_pic->img.plane[2], inBuffer + imgPixels + (imgPixels >> 2), (imgPixels >> 2));   

	m_pic->i_pts = (int64_t)m_frames * m_param->i_fps_den;
	m_pic->i_type = X264_TYPE_AUTO;
	m_pic->i_qpplus1 = 0;

	static x264_picture_t picOut;
	x264_nal_t *nal;
	int nalCnt = 0;
	if (x264_encoder_encode(m_x264, &nal, &nalCnt, m_pic, &picOut) < 0)
	{
		fprintf(stdout, "x264: x264_encoder_encode failed\n");
		return false;
	}
	assert(nalCnt == 1);	// always encode one nal, sps and pps have been encoded in func Open
	if (*outLen < (nal[0].i_payload * 3 >> 1))
	{
		fprintf(stdout, "output buffer is not big enough\n");
		return false;
	}
	x264_nal_encode(outBuffer, outLen, bAnnexB, &nal[0]);
	*frameType = picOut.i_type;

	++m_frames;

#if LOG_INFO
	QueryPerformanceCounter(&litmp);
	end = litmp.QuadPart;
	time = (double)(end - start) / freq;
	m_encodedTime += time;
#endif

	return true;
}

// TODO
//void CX264VideoEncoder::SetResolution( int aWidth, int aHeight )
//{
//	m_width = aWidth;
//	m_height = aHeight;
//}

// must be called after Open()
void CX264VideoEncoder::UpdatePS()
{
	x264_nal_t * nal;
	int nalCnt;
	x264_encoder_headers(m_x264, &nal, &nalCnt);
	x264_nal_encode(m_spsBuffer, &m_spsLen, 0, &nal[1]);
	x264_nal_encode(m_ppsBuffer, &m_ppsLen, 0, &nal[2]);
}

bool CX264VideoEncoder::GetSPS( uint8_t * aSPSBuffer, int * aSPSLen, int aAnnexB )
{
	if (m_spsLen == 0)
		return false;

	assert(aSPSBuffer);
	assert(aSPSLen);
	if (aAnnexB)
	{
		*aSPSLen = sizeof(startCode);
		memcpy(aSPSBuffer, startCode, *aSPSLen);
	}
	else
	{
		*aSPSLen = 0;
	}
	memcpy(aSPSBuffer + *aSPSLen, m_spsBuffer, m_spsLen);
	*aSPSLen += m_spsLen;

	return true;
}

bool CX264VideoEncoder::GetPPS( uint8_t * aPPSBuffer, int * aPPSLen, int aAnnexB )
{
	if (m_ppsLen == 0)
		return false;

	assert(aPPSBuffer);
	assert(aPPSLen);
	if (aAnnexB)
	{
		*aPPSLen = sizeof(startCode);
		memcpy(aPPSBuffer, startCode, *aPPSLen);
	}
	else
	{
		*aPPSLen = 0;
	}
	memcpy(aPPSBuffer + *aPPSLen, m_ppsBuffer, m_ppsLen);
	*aPPSLen += m_ppsLen;

	return true;
}

void CX264VideoEncoder::CustomizeParam()
{
	m_param->i_log_level = X264_LOG_NONE;
	m_param->b_repeat_headers = 0;				// no duplicate sps and pps
	m_param->b_cabac = 0;						// --no-cabac

	m_param->i_width = 352;
	m_param->i_height = 288;

	m_param->i_fps_num = 25;
	m_param->i_fps_den = 1;

	m_param->analyse.b_psnr = 0;				// --no-psnr
	m_param->analyse.b_ssim = 0;				// --no-ssim

	m_param->b_deblocking_filter = 0;			// --no-deblock

	m_param->analyse.i_me_method = X264_ME_DIA;	// --me
	m_param->analyse.i_subpel_refine = 2;		// --subme

	m_param->analyse.inter = 0;					// --partitions 
	m_param->analyse.inter |= X264_ANALYSE_PSUB16x16;	// "p8x8"

	m_param->rc.i_bitrate = 1000;
	m_param->rc.i_rc_method = X264_RC_ABR;

#if LOG_INFO
	m_logFile = fopen("./stat.log", "w");
	if (!m_logFile)
		fprintf(stdout, "Cannot open log file ./stat.log.\n");
#endif
}

void CX264VideoEncoder::LoadConfig(const char * aFile)
{
	FILE * fin = fopen(aFile, "r");
	if (!fin)
	{
		fprintf(stdout, "Cannot open config file %s, using customized params.\n", aFile);
		CustomizeParam();
		return;
	}

	m_param->i_log_level = X264_LOG_NONE;
	m_param->b_repeat_headers = 0;
	m_param->b_cabac = 0;

	char buf[128];
	int bufLen = sizeof(buf);
	memset(buf, 0, bufLen);

	fgets(buf, bufLen, fin);
	sscanf(buf, "width = %d", &m_param->i_width);

	fgets(buf, bufLen, fin);
	sscanf(buf, "height = %d", &m_param->i_height);

	fgets(buf, bufLen, fin);
	sscanf(buf, "fps = %d", &m_param->i_fps_num);
	m_param->i_fps_den = 1;

	fgets(buf, bufLen, fin);
	sscanf(buf, "I interval = %d", &m_param->i_keyint_max);

	int bitrate = 0;
	fgets(buf, bufLen, fin);
	sscanf(buf, "bitrate = %d", &bitrate);
	if (bitrate != 0)
	{
		m_param->rc.i_bitrate = bitrate;
		m_param->rc.i_rc_method = X264_RC_ABR;
	}
		 
	int qp = 0;
	fgets(buf, bufLen, fin);
	sscanf(buf, "qp = %d", &qp);
	if (bitrate == 0 && qp != 0)
	{
		m_param->rc.i_qp_constant = qp;
		m_param->rc.i_rc_method = X264_RC_CQP;
	}

	int b8x8dct = 0;
	fgets(buf, bufLen, fin);
	sscanf(buf, "8x8dct = %d", &b8x8dct);
	m_param->analyse.b_transform_8x8 = b8x8dct;

	char str[64];
	memset(str, 0, sizeof(str));
	fgets(buf, bufLen, fin);
	sscanf(buf, "partitions = %s", str);
	m_param->analyse.inter = 0;
	if (strstr(str, "none"))
	{
		m_param->analyse.inter = 0;
	}
	else if (strstr(str, "all"))
	{
		m_param->analyse.inter = ~0;
	}
	else
	{
		if (strstr(str, "i4x4"))  
			m_param->analyse.inter |= X264_ANALYSE_I4x4;
		if (strstr(str, "i8x8")) 
			m_param->analyse.inter |= X264_ANALYSE_I8x8;
		if (strstr(str, "p8x8")) 
			m_param->analyse.inter |= X264_ANALYSE_PSUB16x16;
		if (strstr(str, "p4x4")) 
			m_param->analyse.inter |= X264_ANALYSE_PSUB8x8;
		if (strstr(str, "b8x8")) 
			m_param->analyse.inter |= X264_ANALYSE_BSUB16x16;
	}

	memset(str, 0, sizeof(str));
	fgets(buf, bufLen, fin);
	sscanf(buf, "me mode = %s", str);
	const char * const ME_MODE[] = {"dia", "hex", "umh", "esa", "tesa"};
	int c = sizeof(ME_MODE) / sizeof(ME_MODE[0]);
	for (int i = 0; i < c; ++i)
	{
		if (strcmp(str, ME_MODE[i]) == 0)
		{
			m_param->analyse.i_me_method = i;
			break;
		}
	}

	fgets(buf, bufLen, fin);
	sscanf(buf, "me range = %d", &m_param->analyse.i_me_range);
	
	fgets(buf, bufLen, fin);
	sscanf(buf, "sub me = %d", &m_param->analyse.i_subpel_refine);

	fgets(buf, bufLen, fin);
	sscanf(buf, "deblock = %d", &m_param->b_deblocking_filter);

	fgets(buf, bufLen, fin);
	sscanf(buf, "psnr = %d", &m_param->analyse.b_psnr);

	fgets(buf, bufLen, fin);
	sscanf(buf, "ssim = %d", &m_param->analyse.b_ssim);

	memset(str, 0, sizeof(str));
	fgets(buf, bufLen, fin);
	sscanf(buf, "log file = %s", str);
#if LOG_INFO
	if (strcmp(str, "") != 0)
	{
		m_logFile = fopen(str, "w");
		if (!m_logFile)
			fprintf(stdout, "Cannot open log file %s.\n", str);
	}
#endif
	
	fclose(fin);

#if 0
	if (m_logFile)
	{
		fprintf(m_logFile, "width = %d\n", m_param->i_width);
		fprintf(m_logFile, "height = %d\n", m_param->i_height);
		fprintf(m_logFile, "fps = %d\n", m_param->i_fps_num / m_param->i_fps_den);
		fprintf(m_logFile, "log level = %d\n", m_param->i_log_level);
		fprintf(m_logFile, "repeat headers = %d\n", m_param->b_repeat_headers);
		fprintf(m_logFile, "psnr = %d\n", m_param->analyse.b_psnr);
		fprintf(m_logFile, "ssim = %d\n", m_param->analyse.b_ssim);
		fprintf(m_logFile, "cabac = %d\n", m_param->b_cabac);
		fprintf(m_logFile, "deblock = %d\n", m_param->b_deblocking_filter);
		fprintf(m_logFile, "me mode = %d\n", m_param->analyse.i_me_method);
		fprintf(m_logFile, "sub me = %d\n", m_param->analyse.i_subpel_refine);
		fprintf(m_logFile, "me range = %d\n", m_param->analyse.i_mv_range);
		fprintf(m_logFile, "inter = %d\n", m_param->analyse.inter);	
		fprintf(m_logFile, "bitrate = %d\n", m_param->rc.i_bitrate);
		fprintf(m_logFile, "rc method = %d\n", m_param->rc.i_rc_method);
		fprintf(m_logFile, "8x8dct = %d\n", m_param->analyse.b_transform_8x8);
		fprintf(m_logFile, "key int = %d\n", m_param->i_keyint_max);
	}
#endif
}
